#ifndef STBIMAGE_HPP
#define STBIMAGE_HPP

#define STB_IMAGE_IMPLEMENTATION
#include "stb_image/stb_image.h"

#define STB_IMAGE_WRITE_IMPLEMENTATION
#include "stb_image/stb_image_write.h"

#include <vector>
#include <Eigen/Eigen>
#include <iostream>

using std::cout;
using std::cerr;
using std::vector;
using std::endl;


class StbImage {
private:
    using uint8 = unsigned char;
    int width, height, n;
    vector<uint8> img;

public:
    StbImage(int w, int h, int n) :width(w), height(h), n(n), img(w * h * n) {}
    StbImage(const std::string& path) {
        uint8* data = stbi_load(path.c_str(), &width, &height, &n, 0);
        if (n != 3) {
            cerr << "Error: Only the images with RGB 3 dimensions were supported." << endl;
            exit(114514);
        }
        else if (!data) {
            cerr << "Error: Image reading failed." << endl;
            exit(114514);
        }
        img = vector<uint8>(data, data + width * height * n);
    }
    uint8 get(int x, int y, int n)const {
        return img[(x + y * width) * this->n + n];
    }
    void set(int x, int y, int n, uint8 val) {
        img[(x + y * width) * this->n + n] = val;
    }
    void setColor(int x, int y, Eigen::Vector3i color) {
        set(x, y, 0, color[0]);
        set(x, y, 1, color[1]);
        set(x, y, 2, color[2]);
    }
    int w()const { return width; }
    int h()const { return height; }
    int nth()const { return n; }
    void write(const std::string& filename) {
        auto p = filename.size();
        --p;
        while (p > 0 && filename[p] != '.') { --p; }
        if (filename[p] != '.') {
            cerr << "Error: Invalid file name" << endl;
            exit(114514);
        }
        auto appendix = filename.substr(p);
        if (appendix == ".png") {
            stbi_write_png(filename.c_str(), width, height, n, img.data(), 0);
        }
        else if (appendix == ".bmp") {
            stbi_write_bmp(filename.c_str(), width, height, n, img.data());
        }else if (appendix == ".jpg") {
            stbi_write_jpg(filename.c_str(), width, height, n, img.data(), 90);
        }
    }
};

#endif //!STBIMAGE_HPP